import os
import uuid
from enum import Enum

from openai import OpenAI

from utils import save_jsonl


class APIMode(Enum):
    CHAT = 'chat'
    COMPLETION = 'completion'
    PREFIX = 'prefix'
    FIM = 'fim'


class DeepSeekClient:
    def __init__(self, api_key, system_prompt=None):
        self.histories = []
        self.api_key = api_key
        self.urls = {
            'chat': os.environ["LLM_URL"]
        }
        self.model = os.environ["LLM_MODEL"]
        self.system_prompt = system_prompt
        self.history = {}

    @staticmethod
    def create_session():
        """创建新会话"""
        return str(uuid.uuid4())

    def generate(self, prompt, mode=APIMode.CHAT, session_id=None, **kwargs):
        """核心调用方法"""
        if mode == APIMode.CHAT:
            if not session_id:
                raise ValueError("session_id is required for chat mode")

            history = self.history.get(session_id, []).copy()

            task = {
                'mode': mode,
                'prompt': prompt,
                'session_id': session_id,
                'history': history,
                'system_prompt': self.system_prompt,
                'kwargs': kwargs
            }

            result = self._work(task)

            if result['success']:
                self.history[session_id] = result['new_history']
                return result['response']
            raise Exception(result['error'])

        else:
            task = {
                'mode': mode,
                'prompt': prompt,
                'system_prompt': self.system_prompt,
                'kwargs': kwargs
            }
            result = self._work(task)
            if result['success']:
                return result['response']
            return ''

    def _work(self, task):
        try:
            mode = task['mode']
            prompt = task['prompt']
            kwargs = task.get('kwargs', {})
            if mode == APIMode.CHAT:
                messages = task['history'].copy()
                system_prompt = self.system_prompt

                # 初始化系统提示
                if not messages and system_prompt:
                    messages.append({'role': 'system', 'content': system_prompt})

                # 添加用户输入
                messages.append({'role': 'user', 'content': prompt})

                # API调用
                assistant_response = self.chat(messages, **kwargs)

                # 更新历史记录
                messages.append({'role': 'assistant', 'content': assistant_response})
                return {
                    'success': True,
                    'response': assistant_response,
                    'new_history': messages
                }
            else:
                method = getattr(self, mode.value, self.completion)
                message = {"role": "user", "content": prompt}
                self.histories.append(message)
                result = method(prompt, **kwargs)
                assistant_message = {"role": "assistant", "content": result}
                self.histories.append(assistant_message)
                return {
                    'success': True,
                    'response': result
                }

        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }

    def chat(self, messages, **kwargs):
        _client = OpenAI(
            api_key=self.api_key,
            base_url=self.urls['chat']
        )
        response = _client.chat.completions.create(
            model=self.model if 'model' not in kwargs else kwargs['model'],
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content

    def completion(self, prompt, **kwargs):
        if self.system_prompt:
            system = {"role": "system", "content": prompt}
        else:
            system = None
        messages = [
            {"role": "user", "content": prompt}
        ]
        if system:
            messages.insert(0, system)
        return self.chat(messages, **kwargs)

    def prefix(self, prompt, language='python', **kwargs):
        _client = OpenAI(
            api_key=self.api_key,
            base_url=self.urls['chat']
        )
        if self.system_prompt:
            system = {"role": "system", "content": prompt}
        else:
            system = None
        messages = [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": f"```{language}\n", "prefix": True}
        ]
        if system:
            messages.insert(0, system)
        response = _client.chat.completions.create(
            model=self.model if 'model' not in kwargs else kwargs['model'],
            messages=messages,
            stop=["```"],
            **kwargs
        )
        return response.choices[0].message.content

    def fim(self, prompt, suffix, max_tokens=128, **kwargs):
        _client = OpenAI(
            api_key=self.api_key,
            base_url=self.urls['chat']
        )
        response = _client.completions.create(
            model=self.model if 'model' not in kwargs else kwargs['model'],
            prompt=prompt,
            suffix=suffix,
            max_tokens=max_tokens,
            **kwargs
        )
        return response.choices[0].text

    def save_histories(self, path):
        save_jsonl(path, self.histories)


# 使用示例
if __name__ == "__main__":
    # 初始化客户端
    client = DeepSeekClient('db8c4c59-9ab0-49f2-9ce6-d0be1d1e7ed2')

    try:
        # 创建新会话
        session_id = client.create_session()

        # 多轮对话
        # print(client.generate("你好，请介绍下你自己", APIMode.CHAT, session_id))
        # print(client.generate("我刚才问了你什么问题？", APIMode.CHAT, session_id))

        # 文本补全
        # print(client.generate("请续写故事：从前有座山", APIMode.COMPLETION))
        # print(client.generate("Please write quick sort code", APIMode.PREFIX, language='python'))
        print(client.generate("def fib(a):", APIMode.FIM, suffix="    return fib(a-1) + fib(a-2)"))
    except Exception as e:
        print(e)
